/****************************************************************
* @file   : cmat.h
* @author : NingJian (freegodly@gmail.com)
* @date   : 06/07/21 16:20:28
* @brief  :  c语言版的矩阵运算实现， 函数对标pytorch 对应的函数
****************************************************************/

#ifndef __CMAT__H__
#define __CMAT__H__

#include <stdio.h>
#include <stdlib.h>

// c语言版的矩阵运算实现， 函数对标pytorch 对应的函数
struct CMat_st
{
    // 数据地址
    double *data_ptr;
    // 数据个数
    size_t data_len;
    // 维度
    int dim;
    // 每个维度的数量记录
    int *dim_size;
};

typedef struct CMat_st CMat;

CMat *CMat_create(int dim, int *dim_size);

CMat *CMat_clone(CMat *mat);

void CMat_delete(CMat *mat);

void CMat_print(CMat *mat);

double CMat_get(CMat *mat, int *dim_info);

void CMat_set(CMat *mat, int *dim_info, double value);

void CMat_view(CMat *mat, int new_dim, int *new_dim_info);

// for this
// torch.arange(0, output_size).unsqueeze(1).repeat(1, output_size)
// torch.arange(0, output_size).unsqueeze(0).repeat(output_size, 1)
CMat *CMat_arange_unsqueeze_repeat(int size, int direction);

CMat *CMat_stack(CMat *x, CMat *y, int dim);

void CMat_unsqueeze(CMat *mat, int dim);

CMat *CMat_repeat(CMat *mat, int *repeat_info);

void CMat_index2diminfo(CMat *mat, size_t index, int *dim_info);

size_t CMat_diminfo2index(CMat *mat, int *dim_info);

void CMat_sigmoid(CMat *mat);

void CMat_exp(CMat *mat);

void CMat_mult_value(CMat *mat, double value);

// 点积是broadcast的。broadcast是torch的一个概念，简单理解就是在一定的规则下允许高维Tensor和低维Tensor之间的运算。broadcast的概念稍显复杂
// 这里举一个点积broadcast的例子。在例子中，a是二维Tensor，b是三维Tensor，但是a的维度与b的后两位相同，那么a和b仍然可以做点积，点积结果是一个和b维度一样的三维Tensor，运算规则是：若c = a * b, 则c[i,*,*] = a * b[i, *, *]，即沿着b的第0维做二维Tensor点积，
void CMat_mult_mat(CMat *mat_a, CMat *mat_b);

void CMat_add_mat(CMat *mat_a, CMat *mat_b);

//将tensor的维度换位
CMat *CMat_permute(CMat *mat, int *dim_info);

CMat *CMat_pick(CMat *mat, int (*pick_dim_info)[2]);

CMat *CMat_cat(CMat **mats, int mats_len, int dim);

#endif //!__CMAT__H__